package com.zhongbo.java.websocket.config;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.text.SimpleDateFormat; // 引入 SimpleDateFormat 类
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import java.util.stream.Stream;

// 配置类，用于启用 WebSocket 功能
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 注册 WebSocket 处理器，指定路径为 "/websocket" 并允许跨域访问
        registry.addHandler(new SimpleChatHandler(), "/websocket")
                .setAllowedOrigins("*");
    }

    // 自定义 WebSocket 处理器，实现聊天功能
    public class SimpleChatHandler extends TextWebSocketHandler {
        private final Set<WebSocketSession> sessions = ConcurrentHashMap.newKeySet(); // 存储所有连接会话
        private final Map<String, WebSocketSession> userSessions = new ConcurrentHashMap<>(); // 用户名与会话映射
        private final Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>(); // Session ID 与会话映射
        private final ObjectMapper objectMapper = new ObjectMapper(); // JSON 解析工具
        private final List<String> messageHistory = Collections.synchronizedList(new ArrayList<>()); // 消息历史记录
        private final Map<String, List<String>> privateMessageHistory = new ConcurrentHashMap<>(); // 私聊消息历史记录

        @Override
        public void afterConnectionEstablished(WebSocketSession session) throws Exception {
            // 获取 URL 中的用户名参数
            String username = null;
            String query = session.getUri().getQuery();
            if (query != null) {
                String[] params = query.split("=");
                if (params.length > 1 && "user".equals(params[0])) {
                    username = params[1];
                }
            }
            // 如果未提供用户名，则生成一个默认用户名
            if (username == null || username.isEmpty()) {
                username = "用户_" + Long.toString(System.currentTimeMillis(), 36);
            }
            sessions.add(session); // 将新会话加入集合
            userSessions.put(username, session); // 将用户名与会话关联
            sessionMap.put(session.getId(), session); // 将 Session ID 与会话关联
            System.out.println("新客户端连接: " + session.getId() + " 用户名: " + username);
            sendUserListToAll(); // 发送在线用户列表给所有客户端
            broadcastMessage(Collections.singletonMap("type", "JOIN"), Collections.singletonMap("user", username)); // 广播用户加入消息
        }

        @Override
        protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
            // 解析接收到的 JSON 消息
            Map<String, String> msg = objectMapper.readValue(message.getPayload(), Map.class);
            ensureMessageFields(msg); // 确保消息字段存在
            String type = msg.get("type");
            String user = msg.get("user");
            String text = msg.get("text");

            // 获取发送消息的用户名
            String senderUsername = userSessions.entrySet().stream()
                    .filter(entry -> entry.getValue().equals(session))
                    .map(Map.Entry::getKey)
                    .findFirst()
                    .orElse("未知用户");

            // 添加发送时间到消息中
            SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss"); // 创建 SimpleDateFormat 实例
            msg.put("timestamp", sdf.format(new Date())); // 使用 SimpleDateFormat 格式化时间

            // 私聊或广播消息处理
            if ("all".equals(user)) {
                msg.put("sender", senderUsername); // 添加发送者用户名
                msg.put("fileType", msg.getOrDefault("fileType", "application/octet-stream")); // 添加文件类型信息
                broadcastMessage(msg); // 广播消息给所有用户
                // 添加公共消息到历史记录
                messageHistory.add(objectMapper.writeValueAsString(msg));
            } else if ("GET_HISTORY".equals(type)) {
                // 处理获取消息历史的请求
                List<String> history = getMessageHistory(senderUsername);
                Map<String, Object> historyResponse = new HashMap<>();
                historyResponse.put("type", "HISTORY");
                historyResponse.put("messages", history);
                // 只发送历史消息给请求的用户
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(historyResponse)));
            } else if ("FILE".equals(type)) {
                // 处理文件消息
                try {
                    msg.put("sender", senderUsername); // 添加发送者用户名
                    msg.put("fileType", msg.getOrDefault("fileType", "application/octet-stream")); // 添加文件类型信息
                    if ("all".equals(user)) {
                        broadcastMessage(msg); // 广播文件消息给所有用户
                        // 添加公共文件消息到历史记录
                        messageHistory.add(objectMapper.writeValueAsString(msg));
                    } else {
                        WebSocketSession targetSession = userSessions.get(user); // 根据用户名查找目标会话
                        if (targetSession == null) {
                            targetSession = sessionMap.get(user); // 根据 Session ID 查找目标会话
                        }
                        if (targetSession != null && targetSession.isOpen()) {
                            // 发送文件消息给接收者和发送者
                            targetSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                            // 记录私聊文件消息历史
                            String privateChatKey = senderUsername + "-" + user;
                            privateMessageHistory.computeIfAbsent(privateChatKey, k -> new ArrayList<>()).add(objectMapper.writeValueAsString(msg));
                        } else {
                            System.out.println("目标用户不在线或不存在: " + user);
                            // 返回错误消息给发送者
                            Map<String, String> errorMessage = new HashMap<>(msg);
                            errorMessage.put("type", "ERROR");
                            errorMessage.put("text", "目标用户不在线或不存在: " + user);
                            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
                        }
                    }
                } catch (Exception e) {
                    System.err.println("Error handling file message: " + e.getMessage());
                    e.printStackTrace();
                    // 返回错误消息给发送者
                    Map<String, String> errorMessage = new HashMap<>();
                    errorMessage.put("type", "ERROR");
                    errorMessage.put("text", "文件发送失败: " + e.getMessage());
                    session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
                }
            } else {
                WebSocketSession targetSession = userSessions.get(user); // 根据用户名查找目标会话
                if (targetSession == null) {
                    targetSession = sessionMap.get(user); // 根据 Session ID 查找目标会话
                }
                if (targetSession != null && targetSession.isOpen()) {
                    msg.put("sender", senderUsername); // 添加发送者用户名
                    msg.put("fileType", msg.getOrDefault("fileType", "application/octet-stream")); // 添加文件类型信息
                    // 发送消息给接收者和发送者
                    targetSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                    session.sendMessage(new TextMessage(objectMapper.writeValueAsString(msg)));
                    // 记录私聊消息历史
                    String privateChatKey = senderUsername + "-" + user;
                    privateMessageHistory.computeIfAbsent(privateChatKey, k -> new ArrayList<>()).add(objectMapper.writeValueAsString(msg));
                } else {
                    System.out.println("目标用户不在线或不存在: " + user);
                    // 返回错误消息给发送者
                    Map<String, String> errorMessage = new HashMap<>(msg);
                    errorMessage.put("type", "ERROR");
                    errorMessage.put("text", "目标用户不在线或不存在: " + user);
                    session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMessage)));
                }
            }
        }

        public List<String> getMessageHistory(String username) {
            List<String> history = new ArrayList<>();

            // 添加公共消息历史
            for (String message : messageHistory) {
                try {
                    Map<String, String> msg = objectMapper.readValue(message, Map.class);
                    ensureMessageFields(msg); // 确保消息字段存在
                    if ("all".equals(msg.get("user"))) {
                        history.add(objectMapper.writeValueAsString(msg));
                    }
                } catch (IOException e) {
                    System.err.println("Error parsing message: " + message);
                    e.printStackTrace();
                }
            }

            // 添加私聊消息历史
            for (Map.Entry<String, List<String>> entry : privateMessageHistory.entrySet()) {
                if (entry.getKey().contains(username)) {
                    for (String msgStr : entry.getValue()) {
                        try {
                            Map<String, String> msg = objectMapper.readValue(msgStr, Map.class);
                            ensureMessageFields(msg); // 确保消息字段存在
                            history.add(objectMapper.writeValueAsString(msg));
                        } catch (IOException e) {
                            System.err.println("Error parsing message: " + msgStr);
                            e.printStackTrace();
                        }
                    }
                }
            }

            // 按时间戳排序历史消息
            history.sort(Comparator.comparing(msg -> {
                try {
                    Map<String, String> messageMap = objectMapper.readValue(msg, Map.class);
                    return new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse(messageMap.get("timestamp"));
                } catch (Exception e) {
                    e.printStackTrace();
                    return new Date(0); // 如果解析失败，返回一个默认时间
                }
            }));

            // 增加日志记录，确保历史消息不为空
            if (history.isEmpty()) {
                System.out.println("No messages found for user: " + username);
            } else {
                System.out.println("Found " + history.size() + " messages for user: " + username);
            }

            return history;
        }


        private void ensureMessageFields(Map<String, String> msg) {
            if (!msg.containsKey("sender")) {
                msg.put("sender", "未知用户");
            }
            if (!msg.containsKey("text")) {
                msg.put("text", "");
            }
            if (!msg.containsKey("timestamp")) {
                msg.put("timestamp", new SimpleDateFormat("yyyy-MM-dd HH:mm:ss").format(new Date()));
            }
            if ("FILE".equals(msg.get("type")) && !msg.containsKey("fileType")) {
                msg.put("fileType", "application/octet-stream");
            }
            // 确保消息类型字段存在
            if (!msg.containsKey("type")) {
                msg.put("type", "UNKNOWN");
            }
            // 确保用户字段存在
            if (!msg.containsKey("user")) {
                msg.put("user", "UNKNOWN");
            }
        }

        @Override
        public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
            // 获取断开连接的用户名
            String username = userSessions.entrySet().stream()
                    .filter(entry -> entry.getValue().equals(session))
                    .map(Map.Entry::getKey)
                    .findFirst()
                    .orElse("未知用户");
            userSessions.values().remove(session); // 移除用户会话
            sessionMap.remove(session.getId()); // 移除 Session ID 映射
            System.out.println("客户端断开连接: " + session.getId());
            sendUserListToAll(); // 发送更新后的在线用户列表
            broadcastMessage(Collections.singletonMap("type", "LEAVE"), Collections.singletonMap("user", username)); // 广播用户离开消息
        }

        // 广播消息给所有连接的客户端
        private void broadcastMessage(Map<String, String> message, Map<String, String> typeInfo) {
            try {
                for (WebSocketSession session : sessions) {
                    if (session.isOpen()) {
                        // 合并消息内容和类型信息
                        Map<String, String> combinedMessage = new HashMap<>(message);
                        combinedMessage.putAll(typeInfo);
                        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(combinedMessage)));
                    }
                }
            } catch (IOException e) {
                System.err.println("Error broadcasting message to session");
                e.printStackTrace();
            }
        }

        private void broadcastMessage(Map<String, String> message) {
            try {
                for (WebSocketSession session : sessions) {
                    if (session.isOpen()) {
                        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
                    }
                }
            } catch (IOException e) {
                System.err.println("Error broadcasting message to session");
                e.printStackTrace();
            }
        }

        // 发送在线用户列表给所有连接的客户端
        private void sendUserListToAll() {
            Set<String> users = userSessions.keySet(); // 获取所有在线用户名
            try {
                for (WebSocketSession session : sessions) {
                    if (session.isOpen()) {
                        // 构建用户列表消息
                        Map<String, Object> userListMessage = new HashMap<>();
                        userListMessage.put("type", "USERS");
                        userListMessage.put("users", users);
                        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(userListMessage)));
                    }
                }
            } catch (IOException e) {
                System.err.println("Error sending user list");
                e.printStackTrace();
            }
        }
    }
}